# -*- coding: utf-8 -*-
"""
@Time 2020/12/14 14:18
@Author 鹄望潇湘
@File Main.py
@Desc 使用FCN_VGG16进行训练的入口主函数
"""

from tools_zsj.Trainer import TrainParameter
from FCN_VGG16_Trainer import FCN_VGG16_Trainer
from PascalDataSet import PascalDataSetGenerator
from FCN_VGG16 import FCN_VGG16
import torch


if __name__ == '__main__':
    original_image_path = "/home/xiaoxiang/DataSet/VOC2012/JPEGImages"
    labeled_image_path = "/home/xiaoxiang/DataSet/VOC2012/SegmentationClass"

    normalized_image_path = "/home/xiaoxiang/DataSet/VOC2012_normalized/JPEGImages"
    normalized_labeled_path = "/home/xiaoxiang/DataSet/VOC2012_normalized/SegmentationClass"

    generator = PascalDataSetGenerator(color_object_dict_path="color_object_dict.pickle")

    """
    
    
    # 生成训练集与测试集
    train_dataset = generator.create_dataset_by_txt(
        "/home/xiaoxiang/DataSet/VOC2012/ImageSets/Segmentation/train.txt",
        original_image_path,
        labeled_image_path)

    val_dataset = generator.create_dataset_by_txt(
        "/home/xiaoxiang/DataSet/VOC2012/ImageSets/Segmentation/val.txt",
        original_image_path,
        labeled_image_path,
        crop_size=(480, 320))
    """

    train_dataset = generator.create_normalized_dataset(
        "/home/xiaoxiang/DataSet/VOC2012_normalized/Segmentation/train.txt",
        normalized_image_path,
        normalized_labeled_path)

    val_dataset = generator.create_normalized_dataset(
        "/home/xiaoxiang/DataSet/VOC2012_normalized/Segmentation/val.txt",
        normalized_image_path,
        normalized_labeled_path)

    # single_dataset = generator.create_dataset_by_image("/home/xiaoxiang/DataSet/VOC2012/JPEGImages/2007_002619.jpg",
                                  #  "/home/xiaoxiang/DataSet/VOC2012/SegmentationClass/2007_002619.png")

    device = torch.device(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
    # 配置训练参数
    parameter = TrainParameter(epoch=1000, batch_size=1, num_works=5, data_shuffle=True, device=device)

    # 创建网络
    net = FCN_VGG16()
    # net.restore_dict("./checkpoint7/checkpoint_best_7.pth")

    # 生成训练器
    trainer = FCN_VGG16_Trainer(fcn_net=net, train_datasets=train_dataset, validate_datasets=val_dataset,
                                parameters=parameter)

    # 开始训练
    trainer.start_train()